import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import MinMaxScaler
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.metrics import classification_report, confusion_matrix, roc_curve, auc
from sklearn.model_selection import GridSearchCV
from sklearn.svm import SVC
from xgboost import XGBClassifier
from datetime import datetime, timedelta

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


class ImprovedQuantTradingSystem:
    """
    改进版量化交易系统类
    """

    def __init__(self, token, stock_list, start_date, end_date):
        """
        初始化
        :param token: Tushare token
        :param stock_list: 股票列表
        :param start_date: 开始日期
        :param end_date: 结束日期
        """
        self.token = token
        self.stock_list = stock_list
        self.start_date = start_date
        self.end_date = end_date
        self.data = None
        self.models = {}
        self.best_model = None

    def get_data(self):
        """
        获取原始数据 - 改进版
        """
        pro = ts.pro_api(self.token)
        stock_data = []

        print("正在下载股票数据...")
        for stock in self.stock_list:
            try:
                # 尝试获取日线数据
                df = pro.daily(ts_code=stock, start_date=self.start_date, end_date=self.end_date)

                # 如果数据为空，尝试其他接口
                if df.empty:
                    print(f"{stock} 日线数据为空，尝试获取周线数据...")
                    df = pro.weekly(ts_code=stock, start_date=self.start_date, end_date=self.end_date)
                    if df.empty:
                        print(f"{stock} 周线数据为空，尝试获取复权数据...")
                        df = pro.pro_bar(ts_code=stock, start_date=self.start_date, end_date=self.end_date, adj='qfq')

                if not df.empty:
                    df['代码'] = stock
                    stock_data.append(df)
                    print(f"已下载 {stock} 数据，共 {len(df)} 条记录")
                else:
                    print(f"无法获取 {stock} 数据")
            except Exception as e:
                print(f"获取 {stock} 数据时出错: {str(e)}")
                continue

        if not stock_data:
            raise ValueError("无法获取任何股票数据，请检查股票代码和时间范围")

        self.data = pd.concat(stock_data)
        self.data['trade_date'] = pd.to_datetime(self.data['trade_date'])
        self.data = self.data.rename(columns={
            'trade_date': '日期', 'close': '收盘', 'high': '最高',
            'low': '最低', 'vol': '成交量', 'amount': '成交额'
        }).sort_values(['代码', '日期'])
        print("\n数据获取完成!")

    def create_features(self):
        """
        改进版特征工程
        """
        print("\n开始特征工程...")

        # 应用特征工程
        self.data = self.data.groupby('代码').apply(self._create_features_for_stock).dropna()

        print("特征工程完成!")

    def _create_features_for_stock(self, df):
        """为单只股票创建特征"""
        # 未来5日收益率标签
        df['未来5日收盘'] = df['收盘'].shift(-5)
        df['未来5日收益率'] = (df['未来5日收盘'] - df['收盘']) / df['收盘']
        df['标签'] = np.where(df['未来5日收益率'] > 0.02, 1, 0)  # 设置2%的阈值

        # 技术指标
        df['5日均线'] = df['收盘'].rolling(5).mean()
        df['10日均线'] = df['收盘'].rolling(10).mean()
        df['20日均线'] = df['收盘'].rolling(20).mean()
        df['60日均线'] = df['收盘'].rolling(60).mean()

        # 均线交叉
        df['5_10_交叉'] = np.where(df['5日均线'] > df['10日均线'], 1, 0)
        df['10_20_交叉'] = np.where(df['10日均线'] > df['20日均线'], 1, 0)

        # RSI
        delta = df['收盘'].diff()
        gain = delta.clip(lower=0)
        loss = -delta.clip(upper=0)
        avg_gain = gain.rolling(14).mean()
        avg_loss = loss.rolling(14).mean()
        rs = avg_gain / avg_loss
        df['RSI'] = 100 - (100 / (1 + rs))

        # MACD
        exp12 = df['收盘'].ewm(span=12, adjust=False).mean()
        exp26 = df['收盘'].ewm(span=26, adjust=False).mean()
        df['MACD'] = exp12 - exp26
        df['MACD信号线'] = df['MACD'].ewm(span=9, adjust=False).mean()
        df['MACD_hist'] = df['MACD'] - df['MACD信号线']

        # 布林带
        df['20日均线'] = df['收盘'].rolling(20).mean()
        df['20日标准差'] = df['收盘'].rolling(20).std()
        df['布林上轨'] = df['20日均线'] + 2 * df['20日标准差']
        df['布林下轨'] = df['20日均线'] - 2 * df['20日标准差']
        df['布林带宽'] = (df['布林上轨'] - df['布林下轨']) / df['20日均线']

        # 波动率和成交量指标
        df['波动率'] = df['收盘'].pct_change().rolling(20).std()
        df['成交量_20日均值'] = df['成交量'].rolling(20).mean()
        df['量价变化率'] = df['成交量'].pct_change() / (df['收盘'].pct_change().abs() + 1e-6)  # 避免除零

        # 动量指标
        df['5日动量'] = df['收盘'].pct_change(5)
        df['10日动量'] = df['收盘'].pct_change(10)
        df['20日动量'] = df['收盘'].pct_change(20)

        # 价格通道
        df['20日最高'] = df['最高'].rolling(20).max()
        df['20日最低'] = df['最低'].rolling(20).min()

        # 价格位置
        df['收盘_20日位置'] = (df['收盘'] - df['20日最低']) / (df['20日最高'] - df['20日最低'] + 1e-6)

        return df

    def preprocess_data(self):
        """
        改进版数据预处理
        """
        print("\n开始数据预处理...")

        # 定义数值型特征
        self.numeric_cols = ['收盘', '5日均线', '10日均线', '20日均线', '60日均线',
                             'RSI', 'MACD', 'MACD_hist', '波动率', '成交量_20日均值',
                             '量价变化率', '5日动量', '10日动量', '20日动量', '收盘_20日位置',
                             '5_10_交叉', '10_20_交叉', '布林带宽']

        # 异常值处理 - 使用更稳健的方法
        for col in self.numeric_cols:
            median = self.data[col].median()
            mad = 1.4826 * np.median(np.abs(self.data[col] - median))  # 稳健标准差估计
            lower_bound = median - 3 * mad
            upper_bound = median + 3 * mad
            self.data[col] = self.data[col].clip(lower_bound, upper_bound)

        # 归一化
        self.scaler = MinMaxScaler()
        self.data[self.numeric_cols] = self.scaler.fit_transform(self.data[self.numeric_cols])

        # 主成分分析
        self.pca = PCA(n_components=0.95)
        self.pca_features = self.pca.fit_transform(self.data[self.numeric_cols])
        self._plot_pca()

        # 相关性分析
        self._plot_correlation()

        # 标签分布分析
        self._plot_label_distribution()

        print("数据预处理完成!")

    def build_models(self):
        """
        改进版模型构建和训练
        """
        print("\n开始构建和训练模型...")

        # 准备数据
        X = self.data[self.numeric_cols]
        y = self.data['标签']
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

        # 定义模型列表
        models = [
            ('随机森林', RandomForestClassifier(n_estimators=200, max_depth=10, min_samples_split=5,
                                                class_weight='balanced', random_state=42)),
            ('梯度提升树', GradientBoostingClassifier(n_estimators=150, learning_rate=0.1,
                                                      max_depth=5, random_state=42)),
            ('XGBoost', XGBClassifier(n_estimators=200, learning_rate=0.1, max_depth=5,
                                      use_label_encoder=False, eval_metric='logloss', random_state=42)),
            ('SVM', SVC(C=1.0, kernel='rbf', probability=True, class_weight='balanced', random_state=42))
        ]

        # 训练和评估模型
        best_score = 0
        for name, model in models:
            print(f"\n训练 {name} 模型...")
            model.fit(X_train, y_train)
            self.models[name] = model

            # 评估模型
            y_pred = model.predict(X_test)
            y_prob = model.predict_proba(X_test)[:, 1]

            print(f"\n{name}模型性能:")
            print(classification_report(y_test, y_pred))

            # 计算AUC
            fpr, tpr, _ = roc_curve(y_test, y_prob)
            roc_auc = auc(fpr, tpr)
            print(f"AUC分数: {roc_auc:.4f}")

            # 选择最佳模型
            if roc_auc > best_score:
                best_score = roc_auc
                self.best_model = model
                self.best_model_name = name

        # 特征重要性分析
        self._plot_feature_importance()

        print(f"\n最佳模型是: {self.best_model_name}, AUC: {best_score:.4f}")
        print("模型训练和评估完成!")

    def mean_reversion_strategy(self, data, stock_code):
        """均值回归策略"""
        stock_data = data[data['代码'] == stock_code].copy()
        stock_data = stock_data.sort_values('日期')

        # 生成交易信号 - 更严格的规则
        stock_data['信号'] = 0
        # 买入信号: 价格低于布林下轨且RSI < 30
        stock_data.loc[(stock_data['收盘'] < stock_data['布林下轨']) &
                       (stock_data['RSI'] < 30), '信号'] = 1
        # 卖出信号: 价格高于布林上轨且RSI > 70
        stock_data.loc[(stock_data['收盘'] > stock_data['布林上轨']) &
                       (stock_data['RSI'] > 70), '信号'] = -1

        return stock_data[['日期', '代码', '收盘', '布林上轨', '布林下轨', 'RSI', '信号']]

    def ml_strategy(self, data, stock_code, model):
        """机器学习策略"""
        stock_data = data[data['代码'] == stock_code].copy()
        stock_data = stock_data.sort_values('日期')

        # 预测概率
        features = [col for col in self.numeric_cols if col in stock_data.columns]
        X = stock_data[features]

        # 确保有足够特征
        if len(features) < 5:
            print(f"股票 {stock_code} 特征不足")
            stock_data['上涨概率'] = 0.5
        else:
            stock_data['上涨概率'] = model.predict_proba(X)[:, 1]

        # 生成交易信号 - 更严格的规则
        stock_data['信号'] = 0
        # 买入信号: 上涨概率高且价格在20日均线以上
        stock_data.loc[(stock_data['上涨概率'] > 0.65) &
                       (stock_data['收盘'] > stock_data['20日均线']), '信号'] = 1
        # 卖出信号: 上涨概率低或价格在20日均线以下
        stock_data.loc[(stock_data['上涨概率'] < 0.35) |
                       (stock_data['收盘'] < stock_data['20日均线']), '信号'] = -1

        return stock_data[['日期', '代码', '收盘', '20日均线', '上涨概率', '信号']]

    def trend_following_strategy(self, data, stock_code):
        """趋势跟踪策略"""
        stock_data = data[data['代码'] == stock_code].copy()
        stock_data = stock_data.sort_values('日期')

        # 生成交易信号
        stock_data['信号'] = 0
        # 买入信号: 短期均线上穿长期均线且成交量放大
        stock_data.loc[(stock_data['5日均线'] > stock_data['20日均线']) &
                       (stock_data['5日均线'].shift(1) <= stock_data['20日均线'].shift(1)) &
                       (stock_data['成交量'] > stock_data['成交量_20日均值']), '信号'] = 1
        # 卖出信号: 短期均线下穿长期均线
        stock_data.loc[(stock_data['5日均线'] < stock_data['20日均线']) &
                       (stock_data['5日均线'].shift(1) >= stock_data['20日均线'].shift(1)), '信号'] = -1

        return stock_data[['日期', '代码', '收盘', '5日均线', '20日均线', '成交量', '成交量_20日均值', '信号']]

    def implement_strategies(self):
        """
        改进版交易策略实现
        """
        print("\n实现交易策略...")

        # 比较策略
        sample_stock = self.stock_list[0]  # 使用第一只股票作为示例
        print(f"\n比较策略在股票 {sample_stock} 上的表现...")

        # 获取策略结果
        mr_results = self.mean_reversion_strategy(self.data, sample_stock)
        ml_results = self.ml_strategy(self.data, sample_stock, self.best_model)
        tf_results = self.trend_following_strategy(self.data, sample_stock)

        # 合并结果
        comparison = mr_results.merge(
            ml_results,
            on=['日期', '代码', '收盘'],
            how='outer',
            suffixes=('_MR', '_ML')
        ).merge(
            tf_results,
            on=['日期', '代码', '收盘'],
            how='outer'
        )
        comparison.rename(columns={'信号': '信号_TF'}, inplace=True)

        # 计算日收益率
        comparison['日收益率'] = comparison['收盘'].pct_change()

        # 计算累计收益
        comparison['MR_累计收益'] = (1 + (comparison['信号_MR'].shift(1) * comparison['日收益率'])).cumprod()
        comparison['ML_累计收益'] = (1 + (comparison['信号_ML'].shift(1) * comparison['日收益率'])).cumprod()
        comparison['TF_累计收益'] = (1 + (comparison['信号_TF'].shift(1) * comparison['日收益率'])).cumprod()
        comparison['BH_累计收益'] = (1 + comparison['日收益率']).cumprod()

        # 处理可能的NaN值
        comparison.fillna(method='ffill', inplace=True)
        comparison.fillna(1, inplace=True)

        # 绘制比较图
        plt.figure(figsize=(14, 7))
        plt.plot(comparison['日期'], comparison['MR_累计收益'], label='均值回归策略')
        plt.plot(comparison['日期'], comparison['ML_累计收益'], label=f'{self.best_model_name}策略')
        plt.plot(comparison['日期'], comparison['TF_累计收益'], label='趋势跟踪策略')
        plt.plot(comparison['日期'], comparison['BH_累计收益'], label='买入持有', linestyle='--')
        plt.title(f'{sample_stock} 不同策略收益比较')
        plt.xlabel('日期')
        plt.ylabel('累计收益率')
        plt.legend()
        plt.grid(True)
        plt.xticks(rotation=45)
        plt.tight_layout()
        plt.savefig(f'{sample_stock}_strategy_comparison.png', bbox_inches='tight')
        plt.show()

        # 打印关键指标
        mr_final = comparison['MR_累计收益'].iloc[-1] - 1
        ml_final = comparison['ML_累计收益'].iloc[-1] - 1
        tf_final = comparison['TF_累计收益'].iloc[-1] - 1
        bh_final = comparison['BH_累计收益'].iloc[-1] - 1

        print(f"\n{sample_stock}策略表现比较:")
        print(f"均值回归策略总收益率: {mr_final:.2%}")
        print(f"{self.best_model_name}策略总收益率: {ml_final:.2%}")
        print(f"趋势跟踪策略总收益率: {tf_final:.2%}")
        print(f"买入持有策略总收益率: {bh_final:.2%}")

        # 保存策略比较结果
        comparison.to_csv(f'{sample_stock}_strategy_comparison.csv', index=False)

    def backtest_system(self, initial_capital=100000, stop_loss=0.03, stop_gain=0.08):
        """
        改进版回测系统
        :param initial_capital: 初始资金
        :param stop_loss: 止损比例 (3%)
        :param stop_gain: 止盈比例 (8%)
        """
        print("\n开始回测交易系统...")

        def backtest_strategy(data, stock_code, signal_col):
            """
            通用策略回测函数
            """
            stock_data = data[data['代码'] == stock_code].copy()
            if len(stock_data) == 0:
                print(f"警告: 股票 {stock_code} 无数据可用")
                return None

            stock_data = stock_data.sort_values('日期')

            # 检查数据有效性
            if pd.isna(stock_data['收盘']).any() or (stock_data['收盘'] <= 0).any():
                print(f"警告: 股票 {stock_code} 包含无效价格数据，已跳过")
                return None

            # 初始化交易状态
            position = 0  # 0:无持仓, 1:持有多头
            entry_price = 0
            capital = initial_capital / len(self.stock_list)  # 均分资金
            shares = 0
            portfolio_value = [capital]  # 初始资金
            trades = []  # 记录交易

            for i in range(1, len(stock_data)):
                current_date = stock_data['日期'].iloc[i]
                current_price = stock_data['收盘'].iloc[i]
                prev_signal = stock_data[signal_col].iloc[i - 1]

                # 检查止损止盈
                if position == 1 and entry_price > 0:
                    ret = (current_price - entry_price) / entry_price
                    if ret <= -stop_loss or ret >= stop_gain:
                        # 平仓
                        close_amount = shares * current_price
                        capital += close_amount
                        profit = close_amount - (shares * entry_price)
                        trades.append({...})
                        shares = 0
                        position = 0
                        portfolio_value.append(capital)
                        continue

                # 执行交易信号
                if prev_signal == 1 and position == 0 and capital > 0:
                    # 新增：检查价格有效性
                    if current_price <= 0 or pd.isna(current_price):
                        continue

                    shares = capital // current_price
                    if shares == 0:  # 资金不足买1股
                        continue

                    cost = shares * current_price
                    if not pd.isna(cost):
                        capital -= cost
                        entry_price = current_price
                        position = 1
                        trades.append({...})

                elif prev_signal == -1 and position == 1:
                    # 卖出
                    close_amount = shares * current_price
                    capital += close_amount
                    profit = close_amount - (shares * entry_price)
                    trades.append({...})
                    shares = 0
                    position = 0

                # 计算当前组合价值
                if position == 1 and not pd.isna(current_price):
                    current_value = capital + (shares * current_price)
                else:
                    current_value = capital
                portfolio_value.append(current_value)

            # 计算策略表现
            if not portfolio_value:
                final_value = initial_capital / len(self.stock_list)
            else:
                final_value = portfolio_value[-1]
            total_return = (final_value / (initial_capital / len(self.stock_list))) - 1

            return {
                '股票代码': stock_code,
                '初始资金': initial_capital / len(self.stock_list),
                '最终价值': final_value,
                '总收益率': total_return,
                '年化收益率': self._annualize_return(total_return, stock_data['日期'].iloc[0],
                                                     stock_data['日期'].iloc[-1]),
                '最大回撤': self._calculate_max_drawdown(portfolio_value),
                '夏普比率': self._calculate_sharpe_ratio(portfolio_value),
                '交易次数': len(trades),
                '交易记录': trades
            }

        # 回测三种策略
        strategies = [
            ('均值回归', '信号'),
            (f'{self.best_model_name}策略', '信号'),
            ('趋势跟踪', '信号')
        ]

        all_results = []
        for strategy_name, signal_col in strategies:
            print(f"\n回测 {strategy_name} 策略...")

            strategy_results = []
            for stock in self.stock_list:
                # 准备数据
                stock_data = self.data[self.data['代码'] == stock].copy()
                if stock_data.empty:
                    print(f"跳过 {stock} - 无数据")
                    continue

                # 生成信号
                if strategy_name == '均值回归':
                    signals = self.mean_reversion_strategy(stock_data, stock)
                elif strategy_name == f'{self.best_model_name}策略':
                    signals = self.ml_strategy(stock_data, stock, self.best_model)
                else:  # 趋势跟踪
                    signals = self.trend_following_strategy(stock_data, stock)

                if signals.empty:
                    print(f"跳过 {stock} - 无有效信号")
                    continue

                # 合并信号到数据
                stock_data = stock_data.merge(
                    signals[['日期', signal_col]],
                    on='日期',
                    how='left'
                ).fillna(0)

                # 回测
                result = backtest_strategy(stock_data, stock, signal_col)
                if result:
                    result['策略'] = strategy_name
                    strategy_results.append(result)
                    print(f"{stock} 回测完成，收益率: {result['总收益率']:.2%}")

            if strategy_results:
                # 计算策略整体表现
                strategy_df = pd.DataFrame(strategy_results)
                total_final_value = strategy_df['最终价值'].sum()
                total_initial_value = strategy_df['初始资金'].sum()
                strategy_return = (total_final_value / total_initial_value) - 1

                print(f"\n{strategy_name}策略整体表现:")
                print(f"总收益率: {strategy_return:.2%}")

                # 保存结果
                strategy_df.to_csv(f'{strategy_name}_backtest_results.csv', index=False)
                all_results.append(strategy_df)

        if not all_results:
            print("警告: 无任何有效的回测结果")
            return pd.DataFrame()

        # 合并所有回测结果
        final_results = pd.concat(all_results)

        # 可视化回测结果
        self._plot_backtest_results(final_results)

        # 保存回测结果
        final_results.to_csv('final_backtest_results.csv', index=False)

        print("\n回测完成!")
        return final_results

    def _annualize_return(self, total_return, start_date, end_date):
        """计算年化收益率"""
        days = (end_date - start_date).days
        if days == 0:
            return 0
        years = days / 365.25
        return (1 + total_return) ** (1 / years) - 1

    def _calculate_sharpe_ratio(self, values, risk_free_rate=0.02):
        """计算夏普比率"""
        returns = np.diff(values) / values[:-1]
        if len(returns) < 2:
            return 0
        excess_returns = returns - risk_free_rate / 252  # 日无风险利率
        return np.mean(excess_returns) / (np.std(excess_returns) + 1e-6) * np.sqrt(252)

    def _calculate_max_drawdown(self, values):
        """计算最大回撤"""
        if not values:
            return 0
        peak = values[0]
        max_dd = 0.0

        for value in values:
            if value > peak:
                peak = value
            dd = (peak - value) / peak
            if dd > max_dd:
                max_dd = dd

        return max_dd

    def _plot_backtest_results(self, results):
        """绘制回测结果"""
        # 各策略表现比较
        strategy_perf = results.groupby('策略').agg({
            '总收益率': 'mean',
            '最大回撤': 'mean',
            '夏普比率': 'mean'
        }).reset_index()

        fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(12, 15))

        # 总收益率
        sns.barplot(x='策略', y='总收益率', data=strategy_perf, ax=ax1)
        ax1.set_title('各策略平均总收益率')
        ax1.set_ylabel('收益率')

        # 最大回撤
        sns.barplot(x='策略', y='最大回撤', data=strategy_perf, ax=ax2)
        ax2.set_title('各策略平均最大回撤')
        ax2.set_ylabel('回撤')

        # 夏普比率
        sns.barplot(x='策略', y='夏普比率', data=strategy_perf, ax=ax3)
        ax3.set_title('各策略平均夏普比率')
        ax3.set_ylabel('夏普比率')

        plt.tight_layout()
        plt.savefig('strategy_performance_comparison.png', bbox_inches='tight')
        plt.show()

        # 各股票表现热图
        pivot_returns = results.pivot(index='股票代码', columns='策略', values='总收益率')
        plt.figure(figsize=(12, 8))
        sns.heatmap(pivot_returns, annot=True, fmt=".2%", cmap="RdYlGn", center=0)
        plt.title('各股票在不同策略下的收益率')
        plt.tight_layout()
        plt.savefig('stock_strategy_performance_heatmap.png', bbox_inches='tight')
        plt.show()

    # 其他辅助方法保持不变...
    def _plot_boxplot(self, title, data, filename):
        """绘制箱线图"""
        plt.figure(figsize=(12, 6))
        sns.boxplot(data=data)
        plt.title(title)
        plt.xticks(rotation=45)
        plt.savefig(filename, bbox_inches='tight')
        plt.show()

    def _plot_pca(self):
        """绘制PCA结果"""
        plt.figure(figsize=(10, 6))
        plt.bar(range(1, self.pca.n_components_ + 1), self.pca.explained_variance_ratio_)
        plt.xlabel('主成分')
        plt.ylabel('解释方差比例')
        plt.title('PCA分析 - 主成分解释方差比例')
        plt.savefig('pca_explained_variance.png', bbox_inches='tight')
        plt.show()

    def _plot_correlation(self):
        """绘制相关性热图"""
        plt.figure(figsize=(16, 14))
        sns.heatmap(self.data[self.numeric_cols].corr(), annot=True, cmap='coolwarm', center=0)
        plt.title('技术指标相关性矩阵')
        plt.savefig('correlation_heatmap.png', bbox_inches='tight')
        plt.show()

    def _plot_label_distribution(self):
        """绘制标签分布"""
        plt.figure(figsize=(6, 4))
        self.data['标签'].value_counts().plot(kind='bar')
        plt.title('标签分布 (1:上涨, 0:下跌)')
        plt.savefig('label_distribution.png', bbox_inches='tight')
        plt.show()

    def _plot_confusion_matrix(self, y_true, y_pred, filename):
        """绘制混淆矩阵"""
        plt.figure(figsize=(6, 6))
        sns.heatmap(confusion_matrix(y_true, y_pred), annot=True, fmt='d', cmap='Blues')
        plt.title('混淆矩阵')
        plt.xlabel('预测标签')
        plt.ylabel('真实标签')
        plt.savefig(filename, bbox_inches='tight')
        plt.show()

    def _plot_roc_curve(self, model, X_test, y_test, filename):
        """绘制ROC曲线"""
        y_prob = model.predict_proba(X_test)[:, 1]
        fpr, tpr, _ = roc_curve(y_test, y_prob)
        roc_auc = auc(fpr, tpr)

        plt.figure(figsize=(8, 6))
        plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC曲线 (AUC = {roc_auc:.2f})')
        plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('假阳性率')
        plt.ylabel('真阳性率')
        plt.title('ROC曲线')
        plt.legend(loc="lower right")
        plt.savefig(filename, bbox_inches='tight')
        plt.show()

    def _plot_feature_importance(self):
        """绘制特征重要性"""
        for name, model in self.models.items():
            if hasattr(model, 'feature_importances_'):
                feature_imp = pd.DataFrame({
                    '特征': self.numeric_cols,
                    '重要性': model.feature_importances_
                }).sort_values('重要性', ascending=False)

                plt.figure(figsize=(12, 8))
                sns.barplot(x='重要性', y='特征', data=feature_imp)
                plt.title(f'{name}特征重要性排序')
                plt.savefig(f'{name}_feature_importance.png', bbox_inches='tight')
                plt.show()


# 主程序
if __name__ == "__main__":
    # 配置参数
    TOKEN = '1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c'  # 替换为你的Tushare token
    STOCK_LIST = ['600519.SH', '000858.SZ', '601318.SH', '600036.SH', '000333.SZ']
    START_DATE = '20210101'
    END_DATE = '20230101'

    # 初始化量化交易系统
    print("初始化量化交易系统...")
    qts = ImprovedQuantTradingSystem(TOKEN, STOCK_LIST, START_DATE, END_DATE)

    # 获取数据
    qts.get_data()

    # 特征工程
    qts.create_features()

    # 数据预处理
    qts.preprocess_data()

    # 构建和训练模型
    qts.build_models()

    # 实现和比较策略
    qts.implement_strategies()

    # 回测交易系统 - 使用更保守的参数
    qts.backtest_system(initial_capital=100000, stop_loss=0.03, stop_gain=0.08)

    print("\n量化交易系统运行完成!")